from openpyxl import load_workbook, Workbook
from gooey import Gooey, GooeyParser
from tkinter import messagebox
import datetime
import configparser
import os
import traceback
from datetime import datetime

# 日志记录模块
import logging
from logging.handlers import RotatingFileHandler

# 设置日志记录器并指定输出格式和文件大小阈值
log_handler = RotatingFileHandler('app.log', maxBytes=1024 * 1024, backupCount=5)
logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s', handlers=[log_handler])


def handle_error(message):
    """处理错误并记录日志。"""
    current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    messagebox.showerror("错误", message)
    logging.error(f"{current_time}: {message}")


def read_excel_columns(file_path, column_indices):
    """
    从给定的 Excel 文件中读取指定列的数据。
    :param file_path: Excel 文件路径。
    :param column_indices: 要读取的列索引列表。
    :return: 包含指定列数据的字典，键为列索引，值为该列数据列表。
    """
    try:
        workbook = load_workbook(filename=file_path)
        sheet = workbook['Sheet1']
        column_data = {}
        for column_index in column_indices:
            def column_generator(sheet, column):
                for row in sheet.iter_rows(values_only=True):
                    cell_value = row[column - 1]
                    if cell_value is not None:
                        yield cell_value
            column_data[column_index] = list(column_generator(sheet, column_index))
        return column_data
    except Exception as e:
        error_message = f"读取 Excel 文件 {file_path} 时出现错误：{str(e)}\n详细错误信息：{traceback.format_exc()}"
        handle_error(error_message)
        return {}


def create_mapping(file_path, column1_index, column2_index):
    """
    从给定的 Excel 文件中创建两列数据的映射关系。
    :param file_path: Excel 文件路径。
    :param column1_index: 第一列索引。
    :param column2_index: 第二列索引。
    :return: 两列数据的映射字典。
    """
    column_data = read_excel_columns(file_path, [column1_index, column2_index])
    if column1_index in column_data and column2_index in column_data:
        return dict(zip(column_data[column1_index], column_data[column2_index]))
    else:
        return {}


def process_file(input_file_path, mapping, output_directory):
    """
    根据映射关系处理输入文件，并将结果保存到指定目录。
    :param input_file_path: 输入文件路径。
    :param mapping: 映射字典。
    :param output_directory: 输出目录路径。
    """
    try:
        workbook = load_workbook(filename=input_file_path)
        workbook_new = Workbook()
        sheet_new = workbook_new.active
        sheet = workbook['Sheet1']
        index_new = 0
        sheet_new.append(["IP", "资产状态", "端口", "协议", "环境", "系统名称", "负责人", "分类"])

        def row_generator(sheet):
            for row_index in range(1, sheet.max_row + 1):
                dd = 'D' + str(row_index)
                port_val = sheet[dd].value
                yield row_index, port_val

        for row_index, port_val in row_generator(sheet):
            if port_val in mapping:
                index_new += 1
                ss = sheet[row_index]
                list_val = [ss[0].value, ss[1].value, ss[3].value, ss[4].value, ss[17].value, ss[18].value, ss[19].value]
                # 根据对应关系获取分类值
                sorting_value = mapping[port_val]
                list_val.append(sorting_value)
                sheet_new.append(list_val)

        # 确保保存路径以反斜杠结尾
        save_path = ensure_directory_path_ends_with_slash(output_directory)
        current_time = datetime.now().strftime("%m%d")
        # 正确构建保存的文件名和路径
        full_save_path = os.path.join(save_path, f"{current_time}资产高危端口统计.xlsx")
        workbook_new.save(full_save_path)
    except Exception as e:
        error_message = f"处理文件时出现错误：{str(e)}\n详细错误信息：{traceback.format_exc()}"
        handle_error(error_message)


# 确保给定的路径字符串以反斜杠 \ 结尾，明确为目录路径
def ensure_directory_path_ends_with_slash(path):
    if not path.endswith('\\'):
        return path + '\\'
    return path


# 配置文件路径
config_file = "config.ini"


def save_config(save_path):
    """
    将保存路径保存到配置文件中。
    """
    try:
        config = configparser.ConfigParser()
        config["DEFAULT"] = {"save_path": save_path}
        with open(config_file, "w") as configfile:
            config.write(configfile)
    except Exception as e:
        error_message = f"保存配置文件时出现错误：{str(e)}\n详细错误信息：{traceback.format_exc()}"
        handle_error(error_message)


def load_config():
    """
    从配置文件中加载上次使用的保存路径。
    """
    try:
        config = configparser.ConfigParser()
        if os.path.exists(config_file):
            config.read(config_file)
            return config["DEFAULT"].get("save_path")
        else:
            return None
    except Exception as e:
        error_message = f"加载配置文件时出现错误：{str(e)}\n详细错误信息：{traceback.format_exc()}"
        handle_error(error_message)
        return None


@Gooey
def main():
    """
    主函数，创建图形用户界面，解析命令行参数，进行数据处理。
    """
    default_path = os.path.join("C:\\Users\\Administrator\\Downloads")
    parser = GooeyParser(description="脚本")
    parser.add_argument("path", help="加载文件", widget="FileChooser", default=default_path)
    saved_path = load_config()
    if saved_path:
        parser.add_argument("--save", help="保存路径", widget="DirChooser", default=saved_path, required=True)
    else:
        parser.add_argument("--save", help="保存路径", widget="DirChooser", required=True)
    global args
    try:
        args = parser.parse_args()
        save_config(args.save)
        print(args.save)
        high_file_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'high.xlsx')
        mapping = create_mapping(high_file_path, 2, 3)
        if mapping:
            process_file(args.path, mapping, args.save)
        else:
            messagebox.showwarning("警告", "读取 high.xlsx 失败，无法继续处理。")
    except Exception as e:
        error_message = f"主程序运行时出现错误：{str(e)}\n详细错误信息：{traceback.format_exc()}"
        handle_error(error_message)


if __name__ == "__main__":
    main()
